﻿using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;

namespace SentenceGeneration
{
    /* I uniquely map terminals and nonterminals to integers, with the start symbol = 0. I call the set of
     * terminals and nonterminals 'symbols'.
     * A production (aka rule) has a nonterminal on the LHS, and maps to a list of 0 or more terminals or nonterminals on the RHS.
     * A nonterminal may be on the LHS of more than one production.
     */
    class Production
    {
        public Production(NonTerminal lhs, List<Symbol> rhs)
        {
            this.lhs = lhs;
            this.rhs = rhs;
        }
        public override string ToString()
        {
            string sRhs = rhs.Select<Symbol, string>(sym => sym.ToString()).Aggregate((working, next) => working + " " + next); ;
            return string.Format("{0} : {1}", lhs.ToString(), sRhs);
        }
        public NonTerminal lhs { get; set; }
        public List<Symbol> rhs { get; set; }
    }

    class Symbol
    {
        public Symbol(string representation)
        {
            this.representation = representation;
        }
        public override string ToString()
        {
            return representation;
        }
        public string representation { get; set; } //must be unique among the set of terminals and nonterminals
    }

    class NonTerminal : Symbol
    {
        public NonTerminal(string representation)
            : base(representation)
        {
            productions = new List<Production>();
        }
        public List<Production> productions { get; set; }
    }

    class Terminal : Symbol
    {
        public Terminal(string representation)
            : base(representation) { }
    }

    public class SentenceGenerator
    {
        enum NonTerminalStatus
        {
            Ready, Unsure, Finished
        }

        //these variables are for bookkeeping purposes
        NonTerminal startSymbol;
        HashSet<Terminal> terminals;
        HashSet<NonTerminal> nonTerminals;
        HashSet<Production> productions;
        Dictionary<Production, NonTerminal> productionsLHS; //what nonterminal is on the LHS of this production?
        Dictionary<string, Symbol> symbolByName; //to assist in building up these other data structures
        Dictionary<Terminal, List<Production>> productionsUsingTerminal;//to assist in building up these other data structures

        //these are filled out by the algorithm itself
        Dictionary<NonTerminal, Production> shortestProduction; //for this nonterminal, which of its rules do we use to generate shortest string? called short in paper.
        Dictionary<NonTerminal, Production> previousProduction; //maps nonterminal to the production that uses it and generates the shortest possible string. called prev in paper
        //TODO revert to Hashtable if you have problems.
        Dictionary<NonTerminal, object> once; //maps a nonterminal to NonTerminalStatus or a Production
        Dictionary<Production, bool> productionUsed; //maps a production to true if it has been used already. called mark in paper
        Stack<Symbol> sentenceInProgress; //stores the symbols as we build up a single output sentence

        public SentenceGenerator()
        {
            startSymbol = null;
            terminals = new HashSet<Terminal>();
            nonTerminals = new HashSet<NonTerminal>();
            productions = new HashSet<Production>();
            symbolByName = new Dictionary<string, Symbol>();
            productionsUsingTerminal = new Dictionary<Terminal, List<Production>>();
            productionsLHS = new Dictionary<Production, NonTerminal>();
            shortestProduction = new Dictionary<NonTerminal, Production>();
            previousProduction = new Dictionary<NonTerminal, Production>();
            once = new Dictionary<NonTerminal, object>();
            productionUsed = new Dictionary<Production, bool>();
            sentenceInProgress = new Stack<Symbol>();
        }

        public void SetStartSymbol(string startSymbol)
        {
            Symbol sym;
            if (!symbolByName.TryGetValue(startSymbol, out sym))
            {
                throw new System.ArgumentException("Start symbol " + startSymbol + " is not a defined terminal.");
            }
            if (sym is Terminal)
            {
                throw new System.ArgumentException("You can't have a terminal for a start symbol.");
            }
            this.startSymbol = (NonTerminal)sym;
        }

        public void AddTerminal(string representation)
        {
            if (!symbolByName.ContainsKey(representation))
            {
                Terminal term = new Terminal(representation);
                symbolByName.Add(representation, term);
                terminals.Add(term);
            }
        }

        public void AddNonTerminal(string representation, List<string> rhs)
        {
            Symbol sym;
            NonTerminal nonterm;
            if (!symbolByName.TryGetValue(representation, out sym) || sym is Terminal)
            {
                nonterm = new NonTerminal(representation);
                //it could be classified as a terminal at the moment. if so, fix
                if (sym is Terminal)
                {
                    terminals.Remove((Terminal)sym);
                    symbolByName.Remove(representation);

                    /* all productions that reference this nonterminal think that it's a 
                     * terminal. update these references */
                    foreach (Production p in productionsUsingTerminal[(Terminal)sym])
                    {
                        for (int i = 0; i < p.rhs.Count; i++)
                        {
                            if (p.rhs[i].representation == representation)
                            {
                                p.rhs[i] = nonterm;
                            }
                        }
                    }
                }
                symbolByName.Add(representation, nonterm);
                nonTerminals.Add(nonterm);
            }
            else
            {
                nonterm = (NonTerminal)sym;
            }

            // now create a production from the List of strings
            List<Symbol> productionRHS = new List<Symbol>();
            Production production = new Production(nonterm, productionRHS);
            foreach (string str in rhs)
            {
                /* we start off assuming everything is a terminal. Until later when
                 * we may call AddNonTerminal with the same string representation and
                 * we fix our mistake.
                 * inefficient, but who cares?
                 */
                if (!symbolByName.TryGetValue(str, out sym))
                {
                    //never heard of this symbol before. assume it's a terminal
                    Terminal term = new Terminal(str);
                    symbolByName.Add(str, term);
                    terminals.Add(term);
                    productionRHS.Add(term);

                    /* also note productions using this terminal, in case we discover later that 
                     * this is really a nonterminal and we need to update all references to it */
                    productionsUsingTerminal.Add(term, new List<Production>());
                    productionsUsingTerminal[term].Add(production);
                }
                else
                {
                    //we've seen this symbol before. don't change it, just add to list
                    productionRHS.Add(sym);
                    if (sym is Terminal)
                    {
                        productionsUsingTerminal[(Terminal)sym].Add(production);
                    }
                }
            }
            nonterm.productions.Add(production);
            bool isNew = productions.Add(production);
            productionsLHS.Add(production, nonterm);
            Debug.Assert(isNew);
        }

        void InitializeStaticTables()
        {
            Dictionary<Production, int> productionShortest = new Dictionary<Production, int>();
            Dictionary<NonTerminal, int> dlen = new Dictionary<NonTerminal, int>();

            //initialize these to infinity
            foreach (Production p in productions)
            {
                productionShortest[p] = int.MaxValue;
            }
            foreach (NonTerminal nt in nonTerminals)
            {
                dlen[nt] = int.MaxValue;
            }

            bool changed;
            do
            {
                changed = FillShortestProduction(productionShortest);
            } while (changed);

            do
            {
                changed = FillPreviousProduction(productionShortest, dlen);
            } while (changed);
        }

        void InitializeDynamicTables()
        {
            foreach (NonTerminal nonterminal in nonTerminals)
            {
                once[nonterminal] = NonTerminalStatus.Ready;
            }
            foreach (Production production in productions)
            {
                productionUsed[production] = false;
            }
        }

        /* this is slen() in the paper
         * returns the length of the shortest string generated from this symbol
         */
        int SymbolsShortestSentence(Dictionary<Production, int> productionShortest, Symbol symbol)
        {
            if (symbol is Terminal)
            {
                //obviously a terminal is a sentence of length 1
                return 1;
            }
            else
            {
                Debug.Assert(symbol is NonTerminal);
            }

            //return minimum rlen for all rules using nonterminal
            int retval = int.MaxValue;
            List<Production> prods = ((NonTerminal)symbol).productions;
            foreach (Production p in prods)
            {
                if (productionShortest[p] < retval)
                {
                    retval = productionShortest[p];
                }
            }
            return retval;
        }

        /* iterate once through the rules, pulling values of productionLength upward. return true if something changed
         * productionLength is the shortest string you can generate using this production. called rlen in paper
         * the real purpose is to fill out shortestProduction
         */
        bool FillShortestProduction(Dictionary<Production, int> productionLength)
        {
            bool changed = false;
            foreach (Production p in productions)
            {
                //TODO why +1 ? why does sum start at 1?
                //maybe because 'dereferencing' a production still costs something computationally?
                int sum = 1;
                foreach (Symbol symbol in p.rhs)
                {
                    if (SymbolsShortestSentence(productionLength, symbol) < int.MaxValue)
                    {
                        sum += SymbolsShortestSentence(productionLength, symbol);
                    }
                    else
                    {
                        //sum is infinite because one of its components is infinite
                        sum = int.MaxValue;
                        break;
                    }
                }

                if (sum < productionLength[p])
                {
                    if (sum < SymbolsShortestSentence(productionLength, p.lhs))
                    {
                        //the shortest sentence for this nonterminal is using this rule
                        shortestProduction[p.lhs] = p;
                    }
                    //this rule has a new shortest sentence
                    productionLength[p] = sum;
                    changed = true;
                }
            }
            return changed;
        }

        /* iterate once through the rules, moving values of dlen downwards. return true if something changed
         * dlen maps a nonterminal to the length of the shortest string that uses it
         * the real purpose is to fill out previousProduction
         */
        bool FillPreviousProduction(Dictionary<Production, int> productionShortest, Dictionary<NonTerminal, int> dlen)
        {
            bool changed = false;
            dlen[startSymbol] = SymbolsShortestSentence(productionShortest, startSymbol);
            foreach (Production p in productions)
            {
                if (dlen[p.lhs] == int.MaxValue)
                {
                    continue;
                }

                //cost is the minimum string length if we use this production
                int cost = dlen[p.lhs] + productionShortest[p] - SymbolsShortestSentence(productionShortest, p.lhs);
                foreach (Symbol symbol in p.rhs)
                {
                    if (symbol is NonTerminal && cost < dlen[(NonTerminal)symbol])
                    {
                        dlen[(NonTerminal)symbol] = cost;
                        previousProduction[(NonTerminal)symbol] = p;
                        changed = true;
                    }
                }
            }
            return changed;
        }

        /* work out paths to unused productions by following previousProduction back up through
         * the grammar, starting at startingNonTerminal
         * TODO I think this is what the short paper meant to say, but double check this anyway
         * TODO is this where my error is?
         */
        void SetupPaths(NonTerminal startingNonTerminal)
        {
            NonTerminal currentNonTerminal = startingNonTerminal;//j=i
            while (currentNonTerminal != startSymbol)
            {
                Production prevProd = previousProduction[currentNonTerminal];//k=prev[j]
                currentNonTerminal = productionsLHS[prevProd]; //j=LHS[k]
                if (once[currentNonTerminal] is Production)
                {
                    //we already know which of this nonterminal's symbol productions we are going to use
                    return;
                }
                else if (sentenceInProgress.Contains(startingNonTerminal))
                {
                    //another paper said this about Unsure: the value of ONCE calculated in the last loop is Unsure
                    /* I think I get this now.
                     * If the starting nonterminal is already on the parse stack, then 
                     * introducing its prevProd may put us into an infinite loop. So we mark 
                     * the previous nonterminal as unsure. Besides, if it's already on the 
                     * stack, why do we need to setup a path to it?
                     */
                    once[currentNonTerminal] = NonTerminalStatus.Unsure;
                }
                else
                {
                    //we want to use the production that will take us to startingNonTerminal (using the shortest path)
                    once[currentNonTerminal] = prevProd;
                    productionUsed[prevProd] = true;
                }
            }
        }

        /* we've decided some non-terminal is ready for a new rule, so find one
         */
        void CalculatePaths()
        {
            /* look through all unused productions. if the LHS nonterminal is ready or unsure,
             * have the nonterminal use this production
             * if some other production has already caused this nonterminal to be used, then
             * we move on and don't use the production
             */
            foreach (Production production in productions)
            {
                if (!productionUsed[production])
                {
                    NonTerminal nonterminal = productionsLHS[production];
                    if (once[nonterminal] is NonTerminalStatus)
                    {
                        NonTerminalStatus lhsStatus = (NonTerminalStatus)once[nonterminal];
                        if (lhsStatus == NonTerminalStatus.Ready || lhsStatus == NonTerminalStatus.Unsure)
                        {
                            //we have an unused production and a nonterminal that hasn't chosen a rule yet.
                            //choose this production for the nonterminal on the LHS
                            once[nonterminal] = production;
                            productionUsed[production] = true;
                            Trace.WriteLine(string.Format("once[{0}] -> {1}", nonterminal.ToString(),
                                production.rhs.Select(val => val.ToString()).Aggregate((current, next)
                                    => current + " " + next)));
                        }
                    }
                }
            }

            //set up a path to arrive at each nonterminal for which we have assigned a production to use
            for (int i = 0; i < once.Count; i++)
            {
                KeyValuePair<NonTerminal, object> kvp = once.ElementAt(i);
                if (kvp.Key == startSymbol)
                {
                    continue;
                }
                else if (kvp.Value is Production)
                {
                    SetupPaths(kvp.Key);
                }
            }

            //if we aren't going to use a nonterminal by now, we must be done with it
            for (int i = 0; i < once.Count; i++)
            {
                KeyValuePair<NonTerminal, object> kvp = once.ElementAt(i);
                if (kvp.Value is NonTerminalStatus && (NonTerminalStatus)kvp.Value == NonTerminalStatus.Ready)
                {
                    once[kvp.Key] = NonTerminalStatus.Finished;
                    Trace.WriteLine(string.Format("marking {0} Finished", kvp.Key));
                }
            }
        }

        //given a nonterminal, choose which rule to use
        Production ChooseProduction(NonTerminal nonterminal)
        {
            Production retval;
            if (once[nonterminal] is NonTerminalStatus)
            {
                NonTerminalStatus nts = (NonTerminalStatus)once[nonterminal];
                if (nts == NonTerminalStatus.Ready || nts == NonTerminalStatus.Unsure)
                {
                    //TODO paper is confusing, it passes calc_paths(N), but their version of calc_paths does not accept an argument N
                    /* I think it works like this: calculate paths finds unused and useful rules 
                     * and sticks them in once[]. After calling CalculatePaths, we look and see 
                     * if once has a production for this nonterminal, and return
                     * that production. If not (meaning there are no downstream unused rules, so 
                     * none of the productions of this nonterminal are useful), 
                     * we return the production that results in the shortest string.
                     */
                    CalculatePaths();
                }
            }

            if (once[nonterminal] is Production)
            {
                retval = (Production)once[nonterminal];

                //set back to ready so that we use this nonterminal again, and choose a new rule next time
                once[nonterminal] = NonTerminalStatus.Ready;
            }
            else
            {
                //calc paths could not give us an unused or useful production for the nonterminal, so take the shortest instead
                retval = shortestProduction[nonterminal];
                productionUsed[retval] = true;
            }
            return retval;
        }

        List<Terminal> GenerateSentence()
        {
            List<Terminal> retval = new List<Terminal>();
            sentenceInProgress.Push(startSymbol);
            while (sentenceInProgress.Count > 0)
            {
                Trace.WriteLine("the stack (from top to bottom)= " + sentenceInProgress.Select(val => val.ToString())
                    .Aggregate((cur, next) => cur + " " + next));
                Symbol symbol = sentenceInProgress.Pop();
                if (symbol is Terminal)
                {
                    retval.Add((Terminal)symbol);
                }
                else
                {
                    Debug.Assert(symbol is NonTerminal);
                    Production production = ChooseProduction((NonTerminal)symbol);
                    List<Symbol> reversed = new List<Symbol>(production.rhs);
                    reversed.Reverse();
                    foreach (Symbol i in reversed)
                    {
                        sentenceInProgress.Push(i);
                    }
                }
            }
            return retval;
        }

        /// <summary>
        /// checks that all terminals and nonterminals (besides the start) are a 
        /// production of some other nonterminal besides themself
        /// Also does a breadth first search to verify that all terminals and nonterminals
        /// are reachable from the start symbol.
        /// </summary>
        /// <returns></returns>
        bool AllSymbolsAreReachable()
        {
            //check nonterminals
            foreach (NonTerminal nonterm in nonTerminals)
            {
                bool isReferencedByOther = false;
                if (nonterm == startSymbol)
                {
                    continue;
                }
                foreach (Production production in productions)
                {
                    if (production.rhs.Contains(nonterm) && production.lhs != nonterm)
                    {
                        isReferencedByOther = true;
                        break;
                    }
                }
                if (!isReferencedByOther)
                {
                    return false;
                }
            }

            //check terminals
            foreach (Terminal term in terminals)
            {
                bool isReferencedByOther = false;
                foreach (Production production in productions)
                {
                    if (production.rhs.Contains(term))
                    {
                        isReferencedByOther = true;
                        break;
                    }
                }
                if (!isReferencedByOther)
                {
                    return false;
                }
            }

            //can a symbol be reached from the start symbol?
            Dictionary<Symbol, bool> symbolIsReachable = new Dictionary<Symbol, bool>();
            Queue<Symbol> bfsQ = new Queue<Symbol>();

            //initialize all symbols to uncreachable
            foreach (Terminal t in terminals)
            {
                symbolIsReachable[t] = false;
            }
            foreach (NonTerminal nt in nonTerminals)
            {
                symbolIsReachable[nt] = false;
            }

            //start symbol is reachable by definition
            symbolIsReachable[startSymbol] = true;
            bfsQ.Enqueue(startSymbol);

            /* as we encounter a symbol in the queue, mark it reachable. add all its children 
             * to the queue, unless we already know this symbol is reachable, in which case 
             * there is no need to add its children to the queue*/
            while (bfsQ.Count > 0)
            {
                Symbol front = bfsQ.Dequeue();
                if (front is Terminal)
                {
                    Terminal t = (Terminal)front;
                    symbolIsReachable[t] = true;
                }
                else
                {
                    Debug.Assert(front is NonTerminal);
                    NonTerminal nt = (NonTerminal)front;
                    symbolIsReachable[front] = true;
                    foreach (Production production in nt.productions)
                    {
                        foreach (Symbol s in production.rhs)
                        {
                            /* if we already know it's reachable, don't add it to the queue, 
                             * because we don't want to process its children a second time */
                            if (!symbolIsReachable[s])
                            {
                                bfsQ.Enqueue(s);
                            }
                        }
                    }
                }
            }

            //if any of the symbols are not reachable, return false
            if (symbolIsReachable.Values.Any(val => !val))
            {
                return false;
            }

            return true;
        }

        public List<List<string>> GenerateSentences()
        {
            InitializeStaticTables();
            InitializeDynamicTables();
            Debug.Assert(AllSymbolsAreReachable());
            bool done;
            List<List<string>> retval = new List<List<string>>();
            do
            {
                List<Terminal> aSentence = GenerateSentence();
                retval.Add(new List<string>(aSentence.Select<Terminal, string>(term => term.representation)));
                done = productionUsed.Values.All(val => val == true);
            } while (!done);
            return retval;
        }
    }
}
